from gmm import *


if __name__ == '__main__':
    train_samples = read_data('TrainSamples.csv')
    train_labels = read_data('TrainLabels.csv')

    train_sample_by_labels = [[] for i in range(10)]
    for sample, label in zip(train_samples, train_labels):
        label = int(label[0])
        train_sample_by_labels[label].append(sample)
    for i in range(10):
        train_sample_by_labels[i] = np.array(train_sample_by_labels[i])

    gmms = [GMM(5, 17) for i in range(10)]
    for i, t in enumerate(zip(gmms, train_sample_by_labels)):
        print('Training GMM'  + str(i))
        gmm, samples = t
        gmm.train(samples, n_iter=5, print_params=False)

    # print()
    # print('Train statistics')
    # num_correct_all_labels = 0
    # num_total_all_labels = 0
    # for label in range(10):
    #     num_correct = 0
    #     samples = train_sample_by_labels[label]
    #     for sample in samples:
    #         scores = [gmm.probability(sample) for gmm in gmms]
    #         pred = np.argmax(scores)
    #         if pred == label:
    #             num_correct += 1
    #     num_total = len(samples)
    #     accuracy = num_correct / num_total
    #     num_correct_all_labels += num_correct
    #     num_total_all_labels += num_total
    #     print('Label ' + str(label) + ', num_correct = ' + str(num_correct) + ', num_total = ' + str(num_total) + ', accuracy = ' + str(accuracy))
    # accuracy = num_correct_all_labels / num_total_all_labels
    # print('All labels, num_correct = ' + str(num_correct_all_labels) + ', num_total = ' + str(num_total_all_labels) + ', accuracy = ' + str(accuracy))

    # Train statistics
    # 1 GMM, n_iter=5, num_correct = 28088, num_total = 30000, accuracy = 0.9362666666666667
    # 2 GMM, n_iter=5, num_correct = 28088, num_total = 30000, accuracy = 0.9362666666666667
    # 3 GMM, n_iter=5, num_correct = 28088, num_total = 30000, accuracy = 0.9362666666666667
    # 4 GMM, n_iter=5, num_correct = 28088, num_total = 30000, accuracy = 0.9362666666666667
    # 5 GMM, n_iter=5, num_correct = 28088, num_total = 30000, accuracy = 0.9362666666666667

    test_samples = read_data('TestSamples.csv')
    test_labels = read_data('TestLabels.csv')
    print()
    print('Test statistics')
    num_correct = 0
    for sample, label in zip(test_samples, test_labels):
        label = int(label[0])
        scores = [gmm.probability(sample) for gmm in gmms]
        pred = np.argmax(scores)
        if pred == label:
            num_correct += 1
    num_total = len(test_samples)
    accuracy = num_correct / num_total
    print('Test, num_correct = ' + str(num_correct) + ', num_total = ' + str(num_total) + ', accuracy = ' + str(accuracy))

    # Test statistics
    # 1 GMM, num_correct = 9332, num_total = 10000, accuracy = 0.9332
    # 2 GMM, num_correct = 9332, num_total = 10000, accuracy = 0.9332
    # 3 GMM, num_correct = 9332, num_total = 10000, accuracy = 0.9332
    # 4 GMM, num_correct = 9332, num_total = 10000, accuracy = 0.9332
    # 5 GMM, num_correct = 9332, num_total = 10000, accuracy = 0.9332
